import torch
from constant import *
import math


def opinion_span_list_2_bio(args, opinion_span_list, max_length, token_length):
    r = [term2id['Invalid']] * max_length
    r[args.sen_pre_len:token_length - 1] = [term2id['O']] * (token_length - 1 - args.sen_pre_len)
    for span in opinion_span_list:
        for index in range(span[0], span[1] + 1):
            if index == span[0]:
                r[args.sen_pre_len + index] = term2id['B']
                pass
            else:
                r[args.sen_pre_len + index] = term2id['I']
                pass
            pass
        pass
    return r


def opinion_span_list_2_se(opinion_span_list, max_length):
    # 这里的 se 指的是 start 和 end 两个列表，其中每个列表代表的是类似于 [0, 0, 0, 1, 0， 0] 这样的列表
    s = [0] * max_length
    e = [0] * max_length
    for span in opinion_span_list:
        s[span[0]] = 1
        e[span[1]] = 1
        pass
    return s, e


bin_str_list_cache = [[]] * 300  # 程序中定义的一个缓存


# 二进制：指数级数据增强判定
# n_bit = 0 时，res = ['0']
# n_bit = 2 时，res = ['00', '01', '10', '11']
# n_bit = 3 时，res = ['000', '001', '010', '011', '100', '101', '110', '111']
def get_bin_str_list_range_n(n_bit: int):
    if n_bit < 0:
        return -1
    if len(bin_str_list_cache[n_bit]) > 0:
        return bin_str_list_cache[n_bit]
    r = []
    x = int(math.pow(2, n_bit))  # 相当于 2**l
    for i in range(x):
        v = f'{i:b}'.rjust(n_bit, '0')
        r.append(v)
        pass
    bin_str_list_cache[n_bit] = r
    return r


def get_extended_one_data():
    pass


# 这个方法把 8 × 62 的数据变成了 85 × 62 的数据
def get_extended_batch_data(args, tokenized_inputs, triplet_list, sentence_token_len_list):
    # tokenized_inputs: {
    # "input_ids": tensor([[], [], []]),
    # "token_type_ids": tensor([[], [], []]),
    # "attention_mask": tensor([[], [], []])
    # }
    current_batch_size, seq_len = tokenized_inputs['input_ids'].size(0), tokenized_inputs['input_ids'].size(1)

    def get_opinion(iii, jjj):
        # args.sen_pre_len 是 tokenizer 序列化 plus_text 之后的长度
        return opinion_span_list_2_se(triplet_list[iii][jjj]['opinion_span_list'], seq_len - args.sen_pre_len)

    inputs_plus_res = {
        'input_ids': tokenized_inputs['input_ids'][0:1].clone(),
        'token_type_ids': tokenized_inputs['token_type_ids'][0:1].clone(),
        'attention_mask': tokenized_inputs['attention_mask'][0:1].clone()
    }  # 这个目前就是当前 batch 当中第一个元素的 input_ids， token_type_ids 和 attention_mask
    label_plus_res = {
        'a_s': torch.tensor([triplet_list[0][0]['aspect_span'][0]]),  # 当前 batch 第一个元素的第一个 aspect 的起始位置
        'a_e': torch.tensor([triplet_list[0][0]['aspect_span'][1]]),  # 当前 batch 第一个元素的第一个 aspect 的终止位置
        'a_s_': torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0),
        'a_e_': torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0),
        'is_on': torch.tensor([1]),
        'o_s': torch.tensor(get_opinion(0, 0)[0]).unsqueeze(0),  # 当前 batch 第一个元素的第一个 opinion 的起始位置，但是形式是 [0, 0, 1, 0, 0, 0, 0] 这种
        'o_e': torch.tensor(get_opinion(0, 0)[1]).unsqueeze(0),  # 当前 batch 第一个元素的第一个 opinion 的终止位置，但是形式是 [0, 0, 1, 0, 0, 0, 0] 这种
        's': torch.tensor([triplet_list[0][0]['sentiment']])  # 情感
    }

    label_plus_res['a_s_'][0][triplet_list[0][0]['aspect_span'][0]] = 1  # 这一段挺无语的，其实就是做了一个上边 get_opinion() 的，把这个对应位置进行了一个标记
    label_plus_res['a_e_'][0][triplet_list[0][0]['aspect_span'][1]] = 1

    inputs_plus_for_test = {
        'input_ids': tokenized_inputs['input_ids'][0:1].clone(),
        'token_type_ids': tokenized_inputs['token_type_ids'][0:1].clone(),
        'attention_mask': tokenized_inputs['attention_mask'][0:1].clone()
    }

    # 上边光考虑了 batch 当中的第一个元素，接下来开始变例整个 batch
    for elem_index in range(current_batch_size):
        triplets_in_current_elem = triplet_list[elem_index]
        # 整个 triplets_in_current_elem 是什么样子的呢？
        # 是一个 list，其中的每个元素是一个 triplet ([aspect_start_i, aspect_end_i], [opinion_start_i, opinion_end_i], sentiment_str)
        t_mask = torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)  # 一维向量变成 [1, n] 的向量
        t_mask[0][:sentence_token_len_list[elem_index]] = 1
        n_triplets_in_current_elem = len(triplets_in_current_elem)
        for triplet_index in range(n_triplets_in_current_elem + 1):
            bin_str_list = get_bin_str_list_range_n(n_triplets_in_current_elem - triplet_index - 1)
            # 产生一个 bin list: ['000', '001', '010', '011', '100', '101', '110', '111']

            if bin_str_list == -1:
                if args.use_context_augmentation == -1:
                    continue
                    pass
                # 更新 token_type_ids
                inputs_plus_res['token_type_ids'] = torch.cat((
                    inputs_plus_res['token_type_ids'],
                    tokenized_inputs['token_type_ids'][elem_index:elem_index + 1].clone()
                ))  # 因为下面这个采用的索引是 elem_index:elem_index + 1 这样的范围，所以可以得到二维的内容，可以直接执行 torch.cat()

                # 更新 input_ids
                inputs_plus_res['input_ids'] = torch.cat((
                    inputs_plus_res['input_ids'],
                    tokenized_inputs['input_ids'][elem_index:elem_index + 1].clone()
                ))

                # 更新 attention_mask
                last_attention_mask = inputs_plus_res['attention_mask'][-1:].clone()
                aspect_span = triplets_in_current_elem[triplet_index - 1]['aspect_span']
                last_attention_mask[0][args.sen_pre_len + aspect_span[0]:args.sen_pre_len + aspect_span[1] + 1] = 0
                # 把上一个 triplet 的 aspect 的 attention 的位置标记为 0
                inputs_plus_res['attention_mask'] = torch.cat((inputs_plus_res['attention_mask'], last_attention_mask))

                # 更新 label
                label_plus_res['a_s'] = torch.cat((label_plus_res['a_s'], torch.tensor([sentiment2id['Invalid']])))
                label_plus_res['a_e'] = torch.cat((label_plus_res['a_e'], torch.tensor([sentiment2id['Invalid']])))
                label_plus_res['a_s_'] = torch.cat((
                    label_plus_res['a_s_'],
                    torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                ))
                label_plus_res['a_e_'] = torch.cat((
                    label_plus_res['a_e_'],
                    torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                ))
                label_plus_res['is_on'] = torch.cat((
                    label_plus_res['is_on'],
                    torch.tensor([0])  # 此时表示并没有检测到任何的方面
                ))
                label_plus_res['o_s'] = torch.cat((
                    label_plus_res['o_s'],
                    torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                ))
                label_plus_res['o_e'] = torch.cat((
                    label_plus_res['o_e'],
                    torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                ))
                label_plus_res['s'] = torch.cat((
                    label_plus_res['s'],
                    torch.tensor([sentiment2id['Invalid']])  # 因为没有方面，就无法得到任何情感的分析
                ))

                continue
                pass

            for bin_str_index in range(len(bin_str_list) - 1, -1, -1):
                if args.use_context_augmentation == 0 and (not (bin_str_list[bin_str_index].count('0') == 0 or bin_str_list[bin_str_index].count('1') == 0)):
                    continue
                if args.use_context_augmentation == -1 and not (triplet_index == 0 and bin_str_list[bin_str_index].count('1') == 0):
                    continue

                inputs_plus_res['token_type_ids'] = torch.cat((  # torch.cat() 是将多个 tensor 从某个一维度进行拼接
                    inputs_plus_res['token_type_ids'],
                    tokenized_inputs['token_type_ids'][elem_index:elem_index + 1].clone()
                ))
                inputs_plus_res['input_ids'] = torch.cat((
                    inputs_plus_res['input_ids'],
                    tokenized_inputs['input_ids'][elem_index:elem_index + 1].clone()
                ))

                current_attention_mask = tokenized_inputs['attention_mask'][elem_index:elem_index + 1].clone()
                if triplet_index != 0:  # 如果不是 triplet 当中的第一个元素
                    for index in range(triplet_index):
                        aspect_start, aspect_end = triplets_in_current_elem[index]['aspect_span']
                        current_attention_mask[0][args.sen_pre_len + aspect_start:args.sen_pre_len + aspect_end + 1] = 0
                        pass
                    pass
                for char_index, char_here in enumerate(bin_str_list[bin_str_index]):
                    if char_here == '1':
                        aspect_start, aspect_end = triplets_in_current_elem[char_index + triplet_index + 1]['aspect_span']
                        current_attention_mask[0][args.sen_pre_len + aspect_start:args.sen_pre_len + aspect_end + 1] = 0
                        pass
                    pass
                inputs_plus_res['attention_mask'] = torch.cat((inputs_plus_res['attention_mask'], current_attention_mask))

                triplet = triplets_in_current_elem[triplet_index]
                label_plus_res['a_s'] = torch.cat((
                    label_plus_res['a_s'],
                    torch.tensor([triplet['aspect_span'][0]])
                ))
                label_plus_res['a_e'] = torch.cat((
                    label_plus_res['a_e'],
                    torch.tensor([triplet['aspect_span'][1]])
                ))
                label_plus_res['is_on'] = torch.cat((
                    label_plus_res['is_on'],
                    torch.tensor([1])
                ))
                label_plus_res['o_s'] = torch.cat((
                    label_plus_res['o_s'],
                    torch.tensor(get_opinion(elem_index, triplet_index)[0]).unsqueeze(0)
                ))
                label_plus_res['o_e'] = torch.cat((
                    label_plus_res['o_e'],
                    torch.tensor(get_opinion(elem_index, triplet_index)[1]).unsqueeze(0)
                ))
                label_plus_res['s'] = torch.cat((
                    label_plus_res['s'],
                    torch.tensor([triplet['sentiment']])
                ))

                t_as = torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                t_as[0][triplet['aspect_span'][0]] = 1
                label_plus_res['a_s_'] = torch.cat((label_plus_res['a_s_'], t_as))

                t_ae = torch.tensor([0] * (seq_len - args.sen_pre_len)).unsqueeze(0)
                t_ae[0][triplet['aspect_span'][1]] = 1
                label_plus_res['a_e_'] = torch.cat((label_plus_res['a_e_'], t_ae))

                pass
            pass
        pass
    return inputs_plus_res, label_plus_res, inputs_plus_for_test


# 测试 get_bin_str_list_range_n 方法
if __name__ == '__main__':
    print(get_bin_str_list_range_n(-4))  # -1
    print(get_bin_str_list_range_n(0))  # ['0']
    print(get_bin_str_list_range_n(1))  # ['0', '1']
    print(get_bin_str_list_range_n(2))  # ['00', '01', '10', '11']
    print(get_bin_str_list_range_n(3))  # ['000', '001', '010', '011', '100', '101', '110', '111']
    pass
